"""
Author: Morphlng
Date: 2024-04-05 18:05:17
LastEditTime: 2025-01-26 15:35:49
LastEditors: Morphlng
Description: Camera image observation
FilePath: /DrivingGym/src/driving_gym/environment/agent/obs/image_obs.py
"""

from __future__ import annotations

import random

import cv2
import gymnasium as gym
import numpy as np

from driving_gym.data.data_provider import Snapshot
from driving_gym.environment.agent.obs.base_obs import BaseObs
from driving_gym.environment.scenario.base_scenario import BaseScenario
from driving_gym.misc.util import override
from driving_gym.simulation.adapter_interface import AdapterInterface


class ImageObs(BaseObs):
    def __init__(self, config: dict, adapter: AdapterInterface):
        super().__init__(config, adapter)

        self.source: str = config["source"]
        if self.source.startswith(self.actor_id + ":"):
            self.source = self.source[len(self.actor_id) + 1 :]

        self.width: int = config.get("width", 168)
        self.height: int = config.get("height", 168)
        self.normalize: bool = config.get("normalize", True)
        self.gray_scale: bool = config.get("gray_scale", False)
        self.crop_type: str = config.get("crop_type", None)

    @override(BaseObs)
    def get_obs(self, snapshot: Snapshot, scenario: BaseScenario) -> np.ndarray:
        """Process the image observation from the snapshot data.

        Args:
            snapshot (Snapshot): One frame of the simulation data

        Returns:
            np.ndarray: Processed image observation, always in (height, width, channel) format (uint8 or float32)
        """
        image = snapshot.data[self.actor_id]["sensors"][self.source]
        h, w = image.shape[:2]

        if self.gray_scale and len(image.shape) == 3 and image.shape[2] == 3:
            image = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)

        if self.crop_type is None:
            image = cv2.resize(image, (self.width, self.height))
            if len(image.shape) == 2:
                image = np.expand_dims(image, axis=-1)
        else:
            if self.crop_type == "center":
                x0 = (w - self.width) // 2
                y0 = (h - self.height) // 2
            elif self.crop_type == "random":
                x0 = random.randint(0, w - self.width)
                y0 = random.randint(0, h - self.height)
            image = image[y0 : y0 + self.height, x0 : x0 + self.width]

        if self.normalize:
            image = image.astype(np.float32) / 255.0
        else:
            image = image.astype(np.uint8)
            if np.max(image) <= 1:
                image *= 255

        return image

    @override(BaseObs)
    def get_observation_space(self) -> gym.spaces.Box:
        if self.normalize:
            low, high = 0, 1
            dtype = np.float32
        else:
            low, high = 0, 255
            dtype = np.uint8

        if self.gray_scale:
            shape = (self.height, self.width, 1)
        else:
            shape = (self.height, self.width, 3)

        return gym.spaces.Box(low=low, high=high, shape=shape, dtype=dtype)
